import torch
import sys

sys.path.append("/Users/zzq/Developer/python/SAM/reproduction/infini-transformer/")

from infini_transformer import InfiniTransformer
from infini_transformer import RoPEEmbeddings

# 实例化RoPEEmbeddings，设置其参数
embedder = RoPEEmbeddings(
    dim=64,  # 必须与InfiniTransformer中的dim_key参数匹配
    seq_len=2048,  # 必须与InfiniTransformer中的segment_len参数匹配
    dim_embedding_pct=0.5,  # 使用的嵌入维度百分比
    base=10000,  # 基数
)

# 创建InfiniTransformer的实例，并设置其参数
tfm = InfiniTransformer(
    dim_input=768,  # 输入的维度
    dim_hidden=2048,  # 隐藏层的维度
    dim_key=64,  # 必须与RoPEEmbeddings中的dim参数匹配
    dim_value=64,  # 值的维度
    num_heads=8,  # 注意力头的数量
    activation="ffngeglu",  # 激活函数类型
    segment_len=2048,  # 必须与RoPEEmbeddings中的seq_len参数匹配
    update="delta",  # 更新策略
    causal=True,  # 是否为因果模型
    positional_embedder=embedder,  # 使用RoPEEmbeddings实例
    init_state_learnable=False,  # 初始状态是否可学习
    dropout=0.1,  # dropout率
)

batch = torch.randn(2, 65536, 768)  # 创建一个随机张量，形状为(2, 65536, 768)
# 其中2是批次大小，65536是序列长度，768是输入维度

output = tfm(batch)  # 将批次数据传入InfiniTransformer实例，得到输出
print(output)
